using System;
using System.Collections.Generic;
using System.Text;

namespace DeflateLib
{
    /// <summary>
    /// Utility class for DEFLATE algorithm.
    /// </summary>
    static class DeflateUtils
    {
        internal const int MaxDistance = 32768;
        internal const int MaxLength = 258;

        internal const int MaxCode = 285;
        internal const int MaxStaticCode = 287;
        internal const int MaxDistanceCode = 29;
        internal const int MaxStaticDistanceCode = 31;
        internal const int MaxHuffmanCodeLength = 15;

        internal const int EndOfBlockCode = 256;

        internal const int NoCompresstionBlockType = 0;
        internal const int StandardCodeBlockType = 1;
        internal const int DynamicCodeBlockType = 2;

        static TreeNode[] staticCodes = new TreeNode[MaxStaticCode + 1];
        static TreeNode[] staticDistanceCodes = new TreeNode[MaxStaticDistanceCode + 1];
        static CodesAndDistances staticCodesAndDistances;

        static int[] lengthsMap = { 16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15 };

        internal static TreeNode[] StaticCodes			
        {
            get { return staticCodes; }
        }

        internal static TreeNode[] StaticDistanceCodes
        {
            get { return staticDistanceCodes; }
        }

        internal static CodesAndDistances StaticCodesAndDistances
        {
            get { return staticCodesAndDistances; }
        }

        internal static int[] LengthsMap
        {
            get { return lengthsMap; }
        }

        static DeflateUtils()
        {
            for (int i = 0; i <= 143; i++)
            {
                staticCodes[i].Code = 0x0030 + i;
                staticCodes[i].Len = 8;
            }
            for (int i = 144; i <= 255; i++)
            {
                staticCodes[i].Code = 0x0190 + i - 144;
                staticCodes[i].Len = 9;
            }
            for (int i = 256; i <= 279; i++)
            {
                staticCodes[i].Code = 0x0000 + i - 256;
                staticCodes[i].Len = 7;
            }
            for (int i = 280; i <= MaxStaticCode; i++)
            {
                staticCodes[i].Code = 0x00C0 + i - 280;
                staticCodes[i].Len = 8;
            }

            for (int i = 0; i <= MaxStaticDistanceCode; i++)
            {
                staticDistanceCodes[i].Code = i;
                staticDistanceCodes[i].Len = 5;
            }

            staticCodesAndDistances = BuildHuffmanTrees(staticCodes, staticDistanceCodes);
        }

        internal static int CalculateLength(int[] counts, int[] codeLengths)
        {
            int sum = 0;
            for (int i = 0; i < counts.Length; i++)
            {
                sum += counts[i] * codeLengths[i];
            }
            return sum;
        }

        internal static int CalculateLengthInStaticCodes(int[] counts, int[] distanceCounts)
        {
            int sum = 0;
            for (int i = 0; i < counts.Length; i++)
            {
                sum += counts[i] * staticCodes[i].Len;
            }

            for (int i = 0; i < distanceCounts.Length; i++)
            {
                sum += distanceCounts[i] * staticDistanceCodes[i].Len;
            }
            return sum;
        }

        internal static TreeNode[] BuildCodes(int[] lengths)
        {
            TreeNode[] tree = new TreeNode[lengths.Length];
            for (int i = 0; i < lengths.Length; i++)
            {
                tree[i].Len = lengths[i];
            }
            BuildCodes(tree);
            return tree;
        }

        internal static void BuildCodes(TreeNode[] tree)
        {
            int maxBits = tree[0].Len;
            for (int i = 1; i < tree.Length; i++)
            {
                if (maxBits < tree[i].Len) maxBits = tree[i].Len;
            }

            int[] bitLengthsCount = new int[maxBits + 1];
            for (int i = 0; i < tree.Length; i++)
            {
                ++bitLengthsCount[tree[i].Len];
            }

            int[] nextCode = new int[maxBits + 1];
            int code = 0;
            bitLengthsCount[0] = 0;
            for (int bits = 1; bits <= maxBits; bits++)
            {
                code = (code + bitLengthsCount[bits - 1]) << 1;
                nextCode[bits] = code;
            }

            for (int n = 0; n < tree.Length; n++)
            {
                int len = tree[n].Len;
                if (len != 0)
                {
                    tree[n].Code = nextCode[len];
                    nextCode[len]++;
                }
            }
        }

        internal static CodesAndDistances BuildHuffmanTrees(TreeNode[] codes, TreeNode[] distances)
        {
            CodesAndDistances result = new CodesAndDistances();
            result.Codes = BuildHuffmanTree(codes);
            result.Distances = BuildHuffmanTree(distances);
            return result;
        }

        internal static HuffmanTree BuildHuffmanTree(TreeNode[] nodes)
        {
            List<IndexedTreeNode> nonEmptyCodes = new List<IndexedTreeNode>();
            for (int i = 0; i < nodes.Length; i++)
            {
                if (nodes[i].Len > 0)
                {
                    IndexedTreeNode node = new IndexedTreeNode();
                    node.Node = nodes[i];
                    node.Index = (ushort)i;
                    nonEmptyCodes.Add(node);
                }
            }
            if (nodes.Length < 1) throw new DeflateLibException("No nodes to build a tree");
            return BuildHuffmanTreeBranch(nonEmptyCodes, 0, 0);           
        }

        private struct IndexedTreeNode
        {
            internal TreeNode Node;
            internal ushort Index;
        }

        private static HuffmanTree BuildHuffmanTreeBranch(List<IndexedTreeNode> nodes, int prefix, int prefixLength)
        {
            List<IndexedTreeNode> zeros = new List<IndexedTreeNode>();
            List<IndexedTreeNode> ones = new List<IndexedTreeNode>();
            HuffmanTree branch = new HuffmanTree();
            branch.IsLeaf = false;
            for (int i = 0; i < nodes.Count; i++)
            {
                if (nodes[i].Node.Len == prefixLength && nodes[i].Node.Code == prefix)
                {
                    branch.IsLeaf = true;
                    branch.SymbolIndex = nodes[i].Index;
                }
                else
                {
                    bool nextBit = ((nodes[i].Node.Code >> (nodes[i].Node.Len - prefixLength - 1)) & 1) != 0;
                    if (nextBit)
                        ones.Add(nodes[i]);
                    else
                        zeros.Add(nodes[i]);
                }
            }
            if (branch.IsLeaf)
            {
                if (ones.Count > 0 || zeros.Count > 0) throw new DeflateLibException("Invalid codes for the tree");
            }
            else
            {
                if (zeros.Count > 0)
                    branch.Zero = BuildHuffmanTreeBranch(zeros, (prefix << 1), prefixLength + 1);
                if (ones.Count > 0)
                    branch.One = BuildHuffmanTreeBranch(ones, (prefix << 1) | 1, prefixLength + 1);
            }
            return branch;
        }

        const int EncodedLengthCodeOffset = 257;

        static int[] encodedLengthStart = { 3,4,5,6,7,8,9,10,
            11,13,15,17,19,23,27,31,35,43,51,59,67,83,99,
            115,131,163,195,227,258};
            
        static int[] encodedLengthAdditionalBits = {0,0,0,0,0,0,0,0,
            1,1,1,1,2,2,2,2,3,3,3,3,4,4,4,4,5,5,5,5,0};

        internal static void DecodeLength(int code, out int start, out int bitsToRead)
        {
            start = encodedLengthStart[code - EncodedLengthCodeOffset];
            bitsToRead = encodedLengthAdditionalBits[code - EncodedLengthCodeOffset];
        }

        internal static void EncodeLength(int length, out int code, out int bits, out int bitsLength)
        {
            int i = Array.BinarySearch<int>(encodedLengthStart, length);
            if (i < 0) i = ~i - 1;

            code = i + EncodedLengthCodeOffset;
            bits = length - encodedLengthStart[i];
            bitsLength = encodedLengthAdditionalBits[i];
        }

        static int[] encodedDistanceStart = { 1,2,3,4, 5,7,9,
            13,17,25, 33,49,65, 97,129,193,257,385,513,769,1025,1537,2049,
            3073,4097,6145,8193,12289,16385,24577 };
        // !notice that 284 shall not have extra = 31

        static int[] encodedDistanceAdditionalBits = { 0,0,0,0,1,1,2,2,3,3,4,4,
            5,5,6,6,7,7,8,8,9,9,10,10,11,11,12,12,13,13 };

        internal static void DecodeDistance(int code, out int start, out int bitsToRead)
        {
            start = encodedDistanceStart[code];
            bitsToRead = encodedDistanceAdditionalBits[code];
        }

        internal static void EncodeDistance(int distance, out int code, out int bits, out int bitsLength)
        {
            int i = Array.BinarySearch<int>(encodedDistanceStart, distance);
            if (i < 0) i = ~i - 1;
            
            code = i;
            bits = distance - encodedDistanceStart[i];
            bitsLength = encodedDistanceAdditionalBits[i];
        }

        internal static int[] GetCodeLengths(int[] counts, int maxCodeLength)
        {
            return PackageMergeAlgorithm.Compute(counts, maxCodeLength);
        }

        internal static int[] GetCodeLengths(int[] counts)
        {
            return PackageMergeAlgorithm.Compute(counts, MaxHuffmanCodeLength);
        }

        internal const int SpecialLengthCodeStart = 16;
        internal const int MaxSpecialLengthCode = 18;
        internal const int RepeatOfSingleLengthCode = 16;
        internal const int RepeatOfZeroLengthCode = 17;
        internal const int RepeatOfZeroLengthLargeCode = 18;

        internal static int GetExtraBitsForLengthCode(int code)
        {
            switch (code)
            {
                case RepeatOfSingleLengthCode:
                    return 2;
                case RepeatOfZeroLengthCode:
                    return 3;
                case RepeatOfZeroLengthLargeCode:
                    return 7;
                default:
                    return 0;
            }
        }

        internal static int[] CompactCodeLengths(int[] counts, int offset, int length)
        {
            List<int> compactCodes = new List<int>();
            int i = 0;
            while(i < length)
            {
                if (counts[offset + i] == 0)
                {
                    int j = 0;
                    do
                    {
                        ++j;
                    } while (i + j < length && j < 138 && counts[offset + i + j] == 0);
                    if (j < 3)
                    {
                        if(j >= 1) compactCodes.Add(0);
                        if (j >= 2) compactCodes.Add(0);
                    }
                    else if (j < 11)
                    {
                        compactCodes.Add(RepeatOfZeroLengthCode);
                        compactCodes.Add(j - 3);
                    }
                    else
                    {
                        compactCodes.Add(RepeatOfZeroLengthLargeCode);
                        compactCodes.Add(j - 11);
                    }
                    i += j;
                }
                else
                {
                    int current = counts[offset + i++];
                    compactCodes.Add(current);
                    int j = 0;
                    while (i + j < length && j < 6 && counts[offset + i + j] == current)
                    {
                        j++;
                    }
                    if (j >= 3)
                    {
                        compactCodes.Add(RepeatOfSingleLengthCode);
                        compactCodes.Add(j - 3);
                        i += j;
                    }
                }
            }
            return compactCodes.ToArray();
        }

        [System.Diagnostics.Conditional("DETAILS")]
        internal static void DumpTrees(string filename, int[] clen, TreeNode[] codes, TreeNode[] distances)
        {
            DumpTrees(filename, clen,
                Array.ConvertAll<TreeNode,int>(codes, delegate(TreeNode n) { return n.Len; }),
                Array.ConvertAll<TreeNode, int>(distances, delegate(TreeNode n) { return n.Len; }));
        }

        [System.Diagnostics.Conditional("DETAILS")]
        internal static void DumpTrees(string filename, int[] clen, int[] codes, int[] distances)
        {
            using (System.IO.StreamWriter sw = new System.IO.StreamWriter(filename, true))
            {
                sw.WriteLine("CLEN:");
                for (int i = 0; i < clen.Length; i++)
                {
                    sw.Write("\t" + clen[i]);
                    if (i % 20 == 9) sw.WriteLine();
                }
                sw.WriteLine();
                sw.WriteLine("CODES:");
                for (int i = 0; i < codes.Length; i++)
                {
                    sw.Write("\t" + codes[i]);
                    if (i % 20 == 9) sw.WriteLine();
                }
                sw.WriteLine();
                sw.WriteLine("DISTANCES:");
                for (int i = 0; i < distances.Length; i++)
                {
                    sw.Write("\t" + distances[i]);
                    if (i % 20 == 9) sw.WriteLine();
                }
                sw.WriteLine();
                sw.WriteLine("----------------------");
            }
        }
    }

    struct TreeNode
    {
        public int Code;
        public int Len;
    }    
}
